from flask import Flask, request, jsonify
from flask_cors import CORS
from PIL import Image
import io
import torch
import torch.nn.functional as F
import torchvision.transforms as T
import pandas as pd
import jwt
import datetime
import sqlite3
import os
from werkzeug.security import generate_password_hash, check_password_hash

# 导入你的模型定义
from model import ResNet_18

# —— 配置 —— #
MODEL_PATH = 'D:/aaa/flask_vue/flaskvue-name/records/models/BaseVGG_.pt'         # 模型权重文件
SIGNNAMES_CSV = 'D:/aaa/flask_vue/flaskvue-name/signnames.csv'      # 标签映射文件
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
SECRET_KEY = 'your-secret-key'  # 在生产环境中使用安全的密钥

# —— 初始化 Flask —— #
app = Flask(__name__)
CORS(app)  # 允许所有跨域请求（生产可配置域名）

# —— 数据库初始化 —— #
def init_db():
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS users (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            username TEXT UNIQUE NOT NULL,
            password TEXT NOT NULL,
            role TEXT NOT NULL DEFAULT 'user'
        )
    ''')
    
    # 创建系统设置表
    c.execute('''
        CREATE TABLE IF NOT EXISTS settings (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            key TEXT UNIQUE NOT NULL,
            value TEXT NOT NULL
        )
    ''')
    
    # 创建预测历史记录表
    c.execute('''
        CREATE TABLE IF NOT EXISTS prediction_history (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id INTEGER NOT NULL,
            image_path TEXT NOT NULL,
            prediction_label TEXT NOT NULL,
            confidence REAL NOT NULL,
            prediction_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users (id)
        )
    ''')
    
    # 创建默认管理员账户
    try:
        c.execute('INSERT INTO users (username, password, role) VALUES (?, ?, ?)',
                 ('admin', generate_password_hash('admin123'), 'admin'))
        c.execute('INSERT INTO users (username, password, role) VALUES (?, ?, ?)',
                 ('111111', generate_password_hash('111111'), 'admin'))
    except sqlite3.IntegrityError:
        pass  # 如果管理员账户已存在，跳过
    
    # 初始化系统设置
    try:
        c.execute('INSERT INTO settings (key, value) VALUES (?, ?)',
                 ('prediction_bg_color', '#ffffff'))
    except sqlite3.IntegrityError:
        pass
    
    conn.commit()
    conn.close()

init_db()

# —— 用户认证相关函数 —— #
def create_token(user_id, username, role):
    payload = {
        'user_id': user_id,
        'username': username,
        'role': role,
        'exp': datetime.datetime.utcnow() + datetime.timedelta(days=1)
    }
    return jwt.encode(payload, SECRET_KEY, algorithm='HS256')

@app.route('/api/register', methods=['POST'])
def register():
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    if not username or not password:
        return jsonify({'message': '用户名和密码不能为空'}), 400

    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    
    try:
        c.execute('INSERT INTO users (username, password) VALUES (?, ?)',
                 (username, generate_password_hash(password)))
        conn.commit()
        return jsonify({'message': '注册成功'}), 201
    except sqlite3.IntegrityError:
        return jsonify({'message': '用户名已存在'}), 400
    finally:
        conn.close()

@app.route('/api/login', methods=['POST'])
def login():
    print("收到登录请求")  # 调试日志
    data = request.get_json()
    username = data.get('username')
    password = data.get('password')

    print(f"登录尝试 - 用户名: {username}")  # 调试日志

    if not username or not password:
        return jsonify({'message': '用户名和密码不能为空'}), 400

    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    
    try:
        # 检查用户是否存在
        c.execute('SELECT id, username, password, role FROM users WHERE username = ?', (username,))
        user = c.fetchone()
        
        print(f"数据库查询结果: {user}")  # 调试日志
        
        if user:
            print(f"找到用户，角色: {user[3]}")  # 调试日志
            if check_password_hash(user[2], password):
                token = create_token(user[0], user[1], user[3])
                response_data = {
                    'token': token,
                    'user': {
                        'username': user[1],
                        'role': user[3]
                    },
                    'message': '登录成功'
                }
                print(f"登录成功，返回数据: {response_data}")  # 调试日志
                return jsonify(response_data)
            else:
                print("密码验证失败")  # 调试日志
                return jsonify({'message': '密码错误'}), 401
        else:
            print("用户不存在")  # 调试日志
            return jsonify({'message': '用户不存在'}), 401
    except Exception as e:
        print(f"登录错误: {str(e)}")  # 调试日志
        return jsonify({'message': f'登录过程中发生错误: {str(e)}'}), 500
    finally:
        conn.close()

# 添加一个测试路由来检查数据库中的用户
@app.route('/api/test/users', methods=['GET'])
def test_users():
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    try:
        c.execute('SELECT id, username, role FROM users')
        users = c.fetchall()
        return jsonify({
            'message': '数据库中的用户列表',
            'users': [{'id': u[0], 'username': u[1], 'role': u[2]} for u in users]
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        conn.close()

# —— 加载模型 —— #
model = ResNet_18()
state = torch.load(MODEL_PATH, map_location=DEVICE)
model.load_state_dict(state)
model.to(DEVICE)
model.eval()

# —— 加载标签映射 —— #
sign_df = pd.read_csv(SIGNNAMES_CSV)
class_map = dict(zip(sign_df.ClassId, sign_df.SignName))

# 添加中文标签映射
chinese_labels = {
    0: "限速20",
    1: "限速30",
    2: "限速50",
    3: "限速60",
    4: "限速70",
    5: "限速80",
    6: "限速结束",
    7: "限速100",
    8: "限速120",
    9: "禁止超车",
    10: "禁止卡车超车",
    11: "十字路口",
    12: "优先道路",
    13: "让行",
    14: "停止",
    15: "禁止车辆通行",
    16: "禁止卡车通行",
    17: "禁止通行",
    18: "注意",
    19: "危险弯道左",
    20: "危险弯道右",
    21: "连续弯道",
    22: "路面不平",
    23: "湿滑路面",
    24: "道路变窄",
    25: "施工",
    26: "交通信号灯",
    27: "注意行人",
    28: "注意儿童",
    29: "注意自行车",
    30: "注意冰雪",
    31: "注意野生动物",
    32: "限速结束",
    33: "直行",
    34: "右转",
    35: "左转",
    36: "直行或右转",
    37: "直行或左转",
    38: "靠右行驶",
    39: "靠左行驶",
    40: "环岛",
    41: "禁止超车结束",
    42: "禁止卡车超车结束"
}

# —— 定义预处理 —— #
def preprocess_image(img):
    # 转换为灰度图
    gray_img = img.convert('L')
    
    # 二值化处理
    threshold = 128
    binary_img = gray_img.point(lambda x: 0 if x < threshold else 255, '1')
    
    # 转回RGB
    binary_rgb = binary_img.convert('RGB')
    
    return binary_rgb

def preprocess_speed_limit(img):
    # 转换为灰度图
    gray_img = img.convert('L')
    
    # 使用自适应阈值进行二值化，更好地处理不同光照条件
    binary_img = gray_img.point(lambda x: 0 if x < 128 else 255, '1')
    
    # 转回RGB
    binary_rgb = binary_img.convert('RGB')
    
    return binary_rgb

# 定义三个不同的预处理流程
transform_color = T.Compose([
    T.Resize((32, 32)),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

transform_binary = T.Compose([
    T.Resize((32, 32)),
    T.ToTensor(),
    T.Normalize(mean=[0.5], std=[0.5])
])

transform_speed = T.Compose([
    T.Resize((32, 32)),
    T.ToTensor(),
    T.Normalize(mean=[0.5], std=[0.5])
])

# —— Predict 接口 —— #
@app.route('/api/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'error': 'no file part'}), 400

    file = request.files['file']
    try:
        # 获取用户ID
        token = request.headers.get('Authorization')
        if not token:
            return jsonify({'error': '未授权访问'}), 401
        
        try:
            payload = jwt.decode(token.split(' ')[1], SECRET_KEY, algorithms=['HS256'])
            user_id = payload['user_id']
        except:
            return jsonify({'error': '无效的token'}), 401

        # 保存上传的图片
        filename = f"{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}_{file.filename}"
        filepath = os.path.join('uploads', filename)
        os.makedirs('uploads', exist_ok=True)
        file.save(filepath)

        # 读取图片
        img = Image.open(filepath).convert('RGB')
        original_size = img.size

        # 保存原始图片尺寸
        target_size = (32, 32)
        
        # 调整图片大小，保持宽高比
        img.thumbnail((max(target_size), max(target_size)), Image.Resampling.LANCZOS)
        
        # 创建新的白色背景图片
        new_img = Image.new('RGB', target_size, (255, 255, 255))
        
        # 将调整后的图片粘贴到中心位置
        paste_x = (target_size[0] - img.size[0]) // 2
        paste_y = (target_size[1] - img.size[1]) // 2
        new_img.paste(img, (paste_x, paste_y))

        # 对图片进行三种预处理
        binary_img = preprocess_image(new_img)
        speed_img = preprocess_speed_limit(new_img)

        # 准备三种输入
        x_color = transform_color(new_img).unsqueeze(0).to(DEVICE)
        x_binary = transform_binary(binary_img).unsqueeze(0).to(DEVICE)
        x_speed = transform_speed(speed_img).unsqueeze(0).to(DEVICE)

        # 进行三次预测
        with torch.no_grad():
            # 彩色图片预测
            out_color = model(x_color)
            prob_color = F.softmax(out_color, dim=1)
            conf_color, idx_color = torch.max(prob_color, 1)

            # 二值化图片预测
            out_binary = model(x_binary)
            prob_binary = F.softmax(out_binary, dim=1)
            conf_binary, idx_binary = torch.max(prob_binary, 1)

            # 限速专用预测
            out_speed = model(x_speed)
            prob_speed = F.softmax(out_speed, dim=1)
            conf_speed, idx_speed = torch.max(prob_speed, 1)

            # 获取预测结果
            class_id_color = int(idx_color.item())
            confidence_color = float(conf_color.item())
            class_id_binary = int(idx_binary.item())
            confidence_binary = float(conf_binary.item())
            class_id_speed = int(idx_speed.item())
            confidence_speed = float(conf_speed.item())

            # 检查是否是限速标志（前9个类别是限速标志）
            is_speed_limit = any(idx.item() < 9 for idx in [idx_color, idx_binary, idx_speed])

            # 如果是限速标志，优先使用限速专用预测结果
            if is_speed_limit:
                final_class_id = class_id_speed
                final_confidence = confidence_speed
                prediction_type = "speed"
            else:
                # 否则选择置信度最高的预测结果
                confidences = [
                    (confidence_color, class_id_color, "color"),
                    (confidence_binary, class_id_binary, "binary"),
                    (confidence_speed, class_id_speed, "speed")
                ]
                final_confidence, final_class_id, prediction_type = max(confidences)

            # 获取前3个最可能的预测结果
            top3_conf, top3_idx = torch.topk(prob_speed if is_speed_limit else prob_color, 3)
            top3_predictions = [
                {
                    'class_id': int(idx.item()),
                    'confidence': float(conf.item()),
                    'label': chinese_labels.get(int(idx.item()), '未知标志'),
                    'english_label': class_map.get(int(idx.item()), 'Unknown')
                }
                for conf, idx in zip(top3_conf[0], top3_idx[0])
            ]

        # 保存预测历史记录
        conn = sqlite3.connect('users.db')
        c = conn.cursor()
        c.execute('''
            INSERT INTO prediction_history 
            (user_id, image_path, prediction_label, confidence)
            VALUES (?, ?, ?, ?)
        ''', (user_id, filepath, chinese_labels.get(final_class_id, '未知标志'), final_confidence))
        conn.commit()
        conn.close()

        return jsonify({
            'class_id': final_class_id,
            'confidence': final_confidence,
            'label': chinese_labels.get(final_class_id, '未知标志'),
            'english_label': class_map.get(final_class_id, 'Unknown'),
            'original_size': original_size,
            'processed_size': target_size,
            'prediction_type': prediction_type,
            'top3_predictions': top3_predictions,
            'color_confidence': confidence_color,
            'binary_confidence': confidence_binary,
            'speed_confidence': confidence_speed,
            'is_speed_limit': is_speed_limit
        })

    except Exception as e:
        return jsonify({'error': f'处理图片时出错: {str(e)}'}), 400

# 添加系统设置相关的API
@app.route('/api/settings', methods=['GET'])
def get_settings():
    # 验证管理员权限
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '未授权访问'}), 401
    
    try:
        payload = jwt.decode(token.split(' ')[1], SECRET_KEY, algorithms=['HS256'])
        if payload['role'] != 'admin':
            return jsonify({'message': '需要管理员权限'}), 403
    except:
        return jsonify({'message': '无效的token'}), 401
    
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('SELECT key, value FROM settings')
    settings = dict(c.fetchall())
    conn.close()
    
    return jsonify(settings)

@app.route('/api/settings', methods=['POST'])
def update_settings():
    # 验证管理员权限
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'message': '未授权访问'}), 401
    
    try:
        payload = jwt.decode(token.split(' ')[1], SECRET_KEY, algorithms=['HS256'])
        if payload['role'] != 'admin':
            return jsonify({'message': '需要管理员权限'}), 403
    except:
        return jsonify({'message': '无效的token'}), 401
    
    data = request.get_json()
    if not data or 'prediction_bg_color' not in data:
        return jsonify({'message': '无效的设置数据'}), 400
    
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('UPDATE settings SET value = ? WHERE key = ?',
             (data['prediction_bg_color'], 'prediction_bg_color'))
    conn.commit()
    conn.close()
    
    return jsonify({'message': '设置已更新'})

# 添加获取当前设置的API（无需管理员权限）
@app.route('/api/settings/current', methods=['GET'])
def get_current_settings():
    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    c.execute('SELECT key, value FROM settings')
    settings = dict(c.fetchall())
    conn.close()
    
    return jsonify(settings)

# 添加获取历史记录的API
@app.route('/api/prediction-history', methods=['GET'])
def get_prediction_history():
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'error': '未授权访问'}), 401
    
    try:
        payload = jwt.decode(token.split(' ')[1], SECRET_KEY, algorithms=['HS256'])
        user_id = payload['user_id']
    except:
        return jsonify({'error': '无效的token'}), 401

    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    try:
        c.execute('''
            SELECT ph.id, ph.image_path, ph.prediction_label, ph.confidence, 
                   ph.prediction_time, u.username
            FROM prediction_history ph
            JOIN users u ON ph.user_id = u.id
            WHERE ph.user_id = ?
            ORDER BY ph.prediction_time DESC
        ''', (user_id,))
        history = c.fetchall()
        
        return jsonify({
            'history': [{
                'id': h[0],
                'image_path': h[1],
                'prediction_label': h[2],
                'confidence': h[3],
                'prediction_time': h[4],
                'username': h[5]
            } for h in history]
        })
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        conn.close()

# 添加删除历史记录的API
@app.route('/api/prediction-history/<int:history_id>', methods=['DELETE'])
def delete_prediction_history(history_id):
    token = request.headers.get('Authorization')
    if not token:
        return jsonify({'error': '未授权访问'}), 401
    
    try:
        payload = jwt.decode(token.split(' ')[1], SECRET_KEY, algorithms=['HS256'])
        user_id = payload['user_id']
    except:
        return jsonify({'error': '无效的token'}), 401

    conn = sqlite3.connect('users.db')
    c = conn.cursor()
    try:
        # 获取图片路径
        c.execute('SELECT image_path FROM prediction_history WHERE id = ? AND user_id = ?',
                 (history_id, user_id))
        result = c.fetchone()
        if not result:
            return jsonify({'error': '记录不存在或无权限删除'}), 404
        
        # 删除图片文件
        try:
            os.remove(result[0])
        except:
            pass  # 忽略文件删除错误
        
        # 删除数据库记录
        c.execute('DELETE FROM prediction_history WHERE id = ? AND user_id = ?',
                 (history_id, user_id))
        conn.commit()
        
        return jsonify({'message': '记录已删除'})
    except Exception as e:
        return jsonify({'error': str(e)}), 500
    finally:
        conn.close()

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True)
